# qa_web.py
from flask import Flask, request, render_template_string
from EmbeddingStorageModule import EmbeddingStorageModule
from PromptEngineerModule import PromptEngineerModule
from agent2 import EnterpriseAgentSystem
from debug import DebugCallbackHandler, TaskChainModule
import markdown  # 用于 Markdown 转 HTML

app = Flask(__name__)

# 初始化模块（只初始化一次）
prompt_engineering = PromptEngineerModule()
callback_handler = DebugCallbackHandler()
task_chain = TaskChainModule(callback_handler=callback_handler)
enterprise_agent = EnterpriseAgentSystem(faiss_index_path="faiss_index")

# 美化版 HTML 前端模板
HTML_TEMPLATE = """
<!doctype html>
<html>
<head>
    <title>低空导航知识库</title>
    <style>
        body {
            font-family: 'Arial', sans-serif;
            background-color: #f5f7fa;
            margin: 0;
            padding: 0;
        }
        .container {
            width: 90%;
            max-width: 800px;
            margin: 50px auto;
            background-color: #fff;
            border-radius: 12px;
            box-shadow: 0 4px 15px rgba(0,0,0,0.1);
            padding: 30px;
        }
        h1 {
            text-align: center;
            color: #2c3e50;
        }
        form {
            display: flex;
            justify-content: center;
            margin-bottom: 30px;
        }
        input[type="text"] {
            width: 70%;
            padding: 12px;
            border: 1px solid #ccc;
            border-radius: 6px;
            font-size: 16px;
        }
        input[type="submit"] {
            padding: 12px 20px;
            margin-left: 10px;
            background-color: #3498db;
            color: #fff;
            border: none;
            border-radius: 6px;
            cursor: pointer;
            font-size: 16px;
        }
        input[type="submit"]:hover {
            background-color: #2980b9;
        }
        .card {
            background-color: #ecf0f1;
            padding: 15px;
            border-radius: 8px;
            margin-bottom: 20px;
            line-height: 1.5;
        }
        .label {
            font-weight: bold;
            color: #34495e;
            margin-bottom: 5px;
        }
        pre {
            background-color: #2d2d2d;
            color: #f8f8f2;
            padding: 10px;
            border-radius: 6px;
            overflow-x: auto;
        }
        code {
            background-color: #f1f1f1;
            padding: 2px 4px;
            border-radius: 4px;
        }
        table {
            border-collapse: collapse;
            width: 100%;
            margin-top: 10px;
        }
        table, th, td {
            border: 1px solid #ccc;
        }
        th, td {
            padding: 8px;
            text-align: left;
        }
    </style>
</head>
<body>
    <div class="container">
        <h1>低空导航知识库</h1>
        <form method="post">
            <input type="text" name="query" placeholder="请输入问题" required>
            <input type="submit" value="提交">
        </form>

        {% if query %}
        <div class="card">
            <div class="label">你的问题：</div>
            <div>{{ query }}</div>
        </div>
        <div class="card">
            <div class="label">任务链结果：</div>
            <div>{{ task_chain_result|safe }}</div>
        </div>
        <div class="card">
            <div class="label">Agent响应：</div>
            <div>{{ agent_response|safe }}</div>
        </div>
        {% endif %}
    </div>
</body>
</html>
"""

@app.route("/", methods=["GET", "POST"])
def index():
    query = None
    task_chain_result = None
    agent_response = None

    if request.method == "POST":
        query = request.form["query"]

        # 生成提示词模板（可选打印）
        prompt = prompt_engineering.create_prompt_template(query=query)

        # 执行任务链
        task_chain_output = task_chain.execute_task_chain(query)
        if isinstance(task_chain_output, dict):
            task_chain_text = task_chain_output.get('final_answer') or task_chain_output.get('summary') or str(task_chain_output)
        else:
            task_chain_text = str(task_chain_output)
        task_chain_result = markdown.markdown(task_chain_text, extensions=['fenced_code', 'tables'])

        # 执行Agent系统
        agent_output = enterprise_agent.execute_query(query, user_react=True)
        if isinstance(agent_output, dict):
            agent_text = agent_output.get('final_answer') or agent_output.get('summary') or str(agent_output)
        else:
            agent_text = str(agent_output)
        agent_response = markdown.markdown(agent_text, extensions=['fenced_code', 'tables'])

        log_file = callback_handler.log_file
        print(f"回调日志存储路径: {log_file}")

    return render_template_string(
        HTML_TEMPLATE,
        query=query,
        task_chain_result=task_chain_result,
        agent_response=agent_response
    )

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=True)
